import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import datasets,losses,layers,metrics,Sequential,optimizers

tf.random.set_seed(22)
np.random.seed(22)

batchsz = 128
# 常见单词的数量
total_words = 10000
max_review_len = 80
embedding_len = 100
(x_train,  y_train), (x_test, y_test) = datasets.imdb.load_data(num_words=total_words)
print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)
# x_train: [b, 80]
# x_test: [b, 80]
x_train = keras.preprocessing.sequence.pad_sequences(x_train, maxlen=max_review_len)
x_test = keras.preprocessing.sequence.pad_sequences(x_test, maxlen=max_review_len)
print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)

db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
db_train = db_train.shuffle(10000).batch(batchsz, drop_remainder=True)
db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))
db_test = db_test.batch(batchsz, drop_remainder=True)

sample = next(iter(db_train))
print('sample:', sample[0].shape, sample[1].shape)


class MyRNN(keras.Model):

    def __init__(self, units):
        super(MyRNN, self).__init__()
        # [b, 64]初始化状态
        self.state0 = [tf.zeros([batchsz, units])]
        self.state1 = [tf.zeros([batchsz, units])]
        # 把text转换成embedding表示
        # [b,80]=>[b, 80, 100]
        self.embedding = layers.Embedding(total_words, embedding_len, input_length=max_review_len)
        # [b,80,100], h_dim:64
        # RNN: cell1, cell2, cell3
        # simpleRNN
        self.rnn_cell0 = layers.SimpleRNNCell(units, )
        self.rnn_cell1 = layers.SimpleRNNCell(units, )
        # fc [b,80,100] => [b,64]
        self.outlayer = layers.Dense(1)

    def call(self, inputs, training=None):
        '''
        net(x) net(x, training=True) :train mode
        net(x, training=False)
        :param inputs: [b, 80]
        :param training:
        :return:
        '''
        # [b,80] => [b, 80, 100]
        x = self.embedding(inputs)
        # [b, 80, 100] => [b,64]
        # ********* 循环神经网络的前向传播
        state0 = self.state0
        for word in tf.unstack(x, axis=1): # 80*[b,100]
            # x * wxh + h * whh
            out, state0 = self.rnn_cell0(word, state0, training)
            # out1, state1 = self.rnn_cell1(out, state1, training)
        # [b, 64] => [b, 1]
        x = self.outlayer(out)
        # 计算出概率
        prob = tf.sigmoid(x)

        return prob


def main():
    units = 64
    epochs = 4

    model = MyRNN(units)
    model.compile(
        optimizer=optimizers.Adam(1e-3),
        loss=tf.losses.BinaryCrossentropy(),
        metrics=['accuracy']
    )
    model.fit(db_train, epochs=epochs, validation_data=db_test)

    model.evaluate(db_test)

    model.save_weights('rnn/rnn_weight.ckpt')


if __name__ == '__main__':
    main()



